# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# PyTorch (https://pytorch.org/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
WithForceCell
"""

from typing import Tuple
import torch
from torch import Tensor
from torch import nn

from .force import WithForceCell
from ...partition import NeighbourList
from ...system import Molecule
from ...potential import ForceCell
from ...sampling.modifier import ForceModifier
from ...function import get_tensor


class WithForceQMMMCell(WithForceCell):
    r"""Module that wraps the simulation system with the atomic force function.

    Args:

        system (Molecule):              Simulation system.

        force (PotentialCell):          Atomic force calculation cell.

        cutoff (float):                 Cutoff distance. Defulat: None

        neighbour_list (NeighbourList): Neighbour list. Default: None

        modifier (ForceModifier):       Force modifier. Default: None

        bias (Bias):                    Bias potential: Default: None

    Supported Platforms:

        ``CPU`` ``GPU``

    Symbols:

        B:  Batchsize, i.e. number of walkers of the simulation.

        A:  Number of the atoms in the simulation system.

        D:  Spatial dimension of the simulation system. Usually is 3.

    """

    def __init__(self,
                 system: Molecule = None,
                 force: ForceCell = None,
                 neighbour_list: NeighbourList = None,
                 modifier: ForceModifier = None,
                 qm_index = None,
                 ):

        super().__init__(system=system,
                         force=force,
                         neighbour_list=neighbour_list,
                         modifier=modifier)
        
        self.qm_index = get_tensor(qm_index)
        if self.qm_index is None:
            self.qm_index = get_tensor(torch.arange(system.coordinate.shape[1]))

    def forward(self,
                energy: Tensor = None,
                ) -> Tuple[Tensor, Tensor, Tensor]:
        r"""Calculate the atomic force of the simulation system.

        Args:
            energy(Tensor): Tensor of shape (B, 1). Data type is float.
                             Total potential energy
                             Default: None

            force(Tensor): Tensor of shape (B, A, D). Data type is float.
                            Atomic forces
                            Default: None

            virial(Tensor): Tensor of shape (B, D, D). Data type is float.
                             System pressure tensor (virial)
                             Default: None

        Return:
            - energy, Tensor of shape (B, 1). Data type is float. Total energy
            - force, Tensor of shape (B, A, D). Data type is float. Atomic force
            - virial, Tensor of shape (B, D, D). Data type is float. System pressure tensor (virial)

        Symbols:
            B:  Batchsize, i.e. number of walkers of the simulation.
            A:  Number of the atoms in the simulation system.
            D:  Spatial dimension of the simulation system. Usually is 3.
        """

        energy_ad = energy
        force_ad = -self.system.coordinate.grad if self.system.coordinate.grad is not None else torch.zeros_like(self.system.coordinate)
        virial_ad = None 
        if self.system.pbc_box is not None and self.system.pbc_box.requires_grad:
           virial_ad = 0.5 * self.system.pbc_box.grad * self.system.pbc_box if self.system.pbc_box.grad is not None else torch.zeros_like(self.system.pbc_box)

        if self.force_function is not None:
            coordinate, pbc_box = self.system.coordinate, self.system.pbc_box
            qm_coordinate = coordinate[:,self.qm_index]
            total_force = torch.zeros_like(force_ad)

            qm_coordinate = qm_coordinate * self.length_unit_scale
            if pbc_box is not None:
                pbc_box = pbc_box * self.length_unit_scale

            neigh_idx = None
            neigh_vec = None
            neigh_dis = None
            neigh_mask = None
            if self.neighbour_list is not None:
                neigh_idx, neigh_mask, neigh_vec, neigh_dis = self.neighbour_list(qm_coordinate, pbc_box)

            energy, force, virial = self.force_function(
                coordinate=qm_coordinate,
                neighbour_index=neigh_idx,
                neighbour_mask=neigh_mask,
                neighbour_vector=neigh_vec,
                neighbour_distance=neigh_dis,
                pbc_box=pbc_box
            )

            energy = energy * self.energy_unit_scale
            force = force * self.force_unit_scale
            if virial is not None:
                virial = virial * self.energy_unit_scale
            total_force[:,self.qm_index] = force

        energy, force, virial = self.force_modifier(energy, energy_ad, total_force, force_ad, virial, virial_ad)
        self.coordinate.grad = -force
        if self.pbc_box is not None and self.pbc_box.grad is not None:
            self.pbc_box.grad = 2 * virial / self.pbc_box

        self._energy.data.copy_(energy)

        return energy
    
